from utils.analysis.common.mediation_common_class import RegressionDetails, RegressionModel, Model4Mediation, \
    Model4Result
import pandas as pd
import statsmodels.api as sm
import pingouin as pg


def cal_f_value(cur_df: pd.DataFrame, x_cols: [], y_col: str):
    f_df = cur_df
    # 计算F值
    cs_name = '常量'  # 初始化常量
    f_df[cs_name] = 1
    # 提取自变量和因变量, 自变量，中介变量和控制变量合并为自变量
    # f_x = f_df[[cs_name, 'JG', "性别", "年龄", "学历", "企业角色"]]
    # f_y = f_df['YY']
    f_x = f_df[x_cols]
    f_y = f_df[y_col]
    # 添加截距项
    f_x = sm.add_constant(f_x)
    # 创建多元线性回归模型
    f_result = sm.OLS(f_y, f_x).fit()
    cur_res = dict()
    cur_res['f'] = f_result.fvalue
    cur_res['p'] = f_result.f_pvalue
    return cur_res


def analysis_model4(cur_df: pd.DataFrame, x, y, m, covar):
    res1 = pg.mediation_analysis(data=cur_df, x=x, m=m, y=y, covar=covar, seed=0, n_boot=5000)

    x1 = [x]
    x1.extend(covar)
    res2 = pg.linear_regression(cur_df[x1], cur_df[y])
    y2x_fp = cal_f_value(cur_df, x1, y)

    x2 = [x, m]
    x2.extend(covar)
    res3 = pg.linear_regression(cur_df[x2], cur_df[y])
    y2xm_fp = cal_f_value(cur_df, x2, y)

    # x3需要的数据和x1一样
    x3 = x1
    res4 = pg.linear_regression(cur_df[x3], cur_df[m])
    m2x_fp = cal_f_value(cur_df, x3, m)

    cur_res = dict()

    cur_res['effect_analysis_table'] = res1  # 效应分析表
    cur_res['regression_y2x'] = res2  # x对y的回归
    cur_res['regression_y2x_fp'] = y2x_fp  # x对y的f值和p值
    cur_res['regression_y2xm'] = res3  # m，x对y的回归
    cur_res['regression_y2xm_fp'] = y2xm_fp  # m，x对y的回归
    cur_res['regression_m2x'] = res4  # x对m的回归
    cur_res['regression_m2x_fp'] = m2x_fp  # x对m的回归
    return cur_res


class MediationModel4:
    def __init__(self, df: pd.DataFrame, x_arr: [], y_arr: [], m_arr: [], covar: [], n_boot=5000):
        self.df = df
        self.x_arr = x_arr
        self.y_arr = y_arr
        self.m_arr = m_arr
        self.covar = covar
        self.n_boot = n_boot

    def item_analysis(self, x, y, m) -> Model4Result:
        cur_df = self.df
        res1 = pg.mediation_analysis(data=cur_df, x=x, m=m, y=y, covar=self.covar, seed=0, n_boot=self.n_boot)

        x1 = [x]
        x1.extend(self.covar)
        res2 = pg.linear_regression(cur_df[x1], cur_df[y])
        y2x_fp = cal_f_value(cur_df, x1, y)

        x2 = [x, m]
        x2.extend(self.covar)
        res3 = pg.linear_regression(cur_df[x2], cur_df[y])
        y2xm_fp = cal_f_value(cur_df, x2, y)

        # x3需要的数据和x1一样
        x3 = x1
        res4 = pg.linear_regression(cur_df[x3], cur_df[m])
        m2x_fp = cal_f_value(cur_df, x3, m)

        # 表
        mediation_table = [Model4Mediation(*row) for row in res1.to_numpy()]
        # 模型1
        details_y2x = [RegressionDetails(*row) for row in res2.to_numpy()]
        model_y2x = RegressionModel(f=y2x_fp.get('f'), p=y2x_fp.get('p'), details=details_y2x)
        # 模型2
        details_y2xm = [RegressionDetails(*row) for row in res3.to_numpy()]
        model_y2xm = RegressionModel(f=y2xm_fp.get('f'), p=y2xm_fp.get('p'), details=details_y2xm)
        # 模型3
        details_m2x = [RegressionDetails(*row) for row in res4.to_numpy()]
        model_m2x = RegressionModel(f=m2x_fp.get('f'), p=m2x_fp.get('p'), details=details_m2x)
        # 总返回结果
        res_obj = Model4Result(mediation_table=mediation_table, r_y2x=model_y2x, r_y2xm=model_y2xm, r_m2x=model_m2x,
                               x=x, y=y, m=m, covar=self.covar)
        return res_obj

    def analysis(self):
        res = []
        for x in self.x_arr:
            for y in self.y_arr:
                for m in self.m_arr:
                    item_res = self.item_analysis(x, y, m)
                    res.append(item_res)
        return res


if __name__ == '__main__':
    pd.set_option('expand_frame_repr', False)
    df = pd.read_excel('./test_datas.xlsx')
    df = df[["性别", "年龄", "学历", "企业角色", 'JG', 'YY', 'JX', 'XW']]
    # 1.此文件的结果不包含整体的F值和p值
    # 2.若有多个x,y,m，则需要遍历进行获取结果
    x_arr = ['JG']
    y_arr = ['YY']
    m_arr = ['XW']
    covar = ["性别", "年龄", "学历", "企业角色"]
    res = []
    for x in x_arr:
        for y in y_arr:
            for m in m_arr:
                item_res = analysis_model4(df, x, y, m, covar)
                res.append(item_res)
                for cc in item_res:
                    print(cc)
                    print(item_res[cc])
                    o = item_res[cc]
                    for c_o in o:
                        if isinstance(c_o, Model4Mediation):
                            print(c_o.se)
    for i in res:
        print(i)
    obj = MediationModel4(df, x_arr, y_arr, m_arr, covar)
    result = obj.analysis()
    print(result)
    print('===================================================================================================')
